package class14;

import java.util.*;

/**
 * 实现并查集数据结构
 */
public class Code05_UnionFind {

    // 封装数据
    public static class Node<T> {
        public T value;

        public Node(T value) {
            this.value = value;
        }
    }

    // 并查集数据结构
    public static class UnionFind<T> {
        Map<T, Node<T>> nodes;  // 将所有结点封装一层，让所有数据都是引用类型
        Map<Node<T>, Node<T>> parent;  // 记录每个结点的直系父节点
        Map<Node<T>, Integer> nodeSize;  // 记录每个以头为结点的集合大小

        // 初始化所有容器
        public UnionFind(List<T> list) {
            nodes = new HashMap<>();
            parent = new HashMap<>();
            nodeSize = new HashMap<>();
            for (T li : list) {
                Node<T> node = new Node<>(li);  // 封装数据
                nodes.put(li, node);  // 把
                parent.put(node, node);
                nodeSize.put(node, 1);
            }
        }

        // 找到结点的直系父节点（通过将所有集合扁平化，子节点可以直接找到头结点，优化查找）
        public Node<T> findFather(Node<T> node) {
            Stack<Node<T>> stack = new Stack<>();
            while (node != parent.get(node)) {
                stack.push(node);
                node = parent.get(node);
            }
            while (!stack.isEmpty()) {
                parent.put(stack.pop(), node);
            }
            return node;
        }

        // 判断两个结点是否属于同一集合
        public boolean isSameSet(T a, T b) {
            return findFather(nodes.get(a)) == findFather(nodes.get(b));
        }

        // 合并两个结点或集合
        public void union(T a, T b) {
            Node<T> aHead = findFather(nodes.get(a));
            Node<T> bHead = findFather(nodes.get(b));
            // 如果两个结点的头结点相同，说明在同一集合中，不需要合并，直接返回
            if (aHead == bHead) {
                return;
            }
            // 不相同，就由集合结点少的向集合结点多的合并
            int aSetSize = nodeSize.get(aHead);
            int bSetSize = nodeSize.get(bHead);
            Node<T> bigNode = aSetSize >= bSetSize ? aHead : bHead;
            Node<T> smallNode = aHead == bigNode ? bHead : aHead;
            // 找到谁是元素结点多是那个开始合并
            parent.put(smallNode, bigNode);
            nodeSize.put(bigNode, aSetSize + bSetSize);
            // 删除被合并的集合头结点
            nodeSize.remove(smallNode);
        }
    }

}
